# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/18 10:48 
@Author  : ZhangShenao 
@File    : 3.使用CacheBackedEmbeddings缓存Embedding结果.py 
@Desc    : 使用CacheBackedEmbeddings缓存Embedding结果

CacheBackedEmbeddings是一个支持缓存的Embedding包装器
它可以将Embedding的结果缓存在键值存储中
具体操作是: 对文本进行哈希处理,并将此哈希值用作缓存的键

通过缓存,可以有效提升Embedding的执行效率,减少token的消耗

"""
import os

import dotenv
from langchain.embeddings import CacheBackedEmbeddings
from langchain.storage import LocalFileStore
from langchain_openai import OpenAIEmbeddings

# 加载环境变量
dotenv.load_dotenv()

# 创建底层的Embedding模型
underlying_embedding = OpenAIEmbeddings(
    model='text-embedding-3-small',
    dimensions=1024,
    openai_api_base=os.getenv('OPENAI_API_BASE')
)

# 创建LocalFileStore,在本地文件中存储Embedding结果
store = LocalFileStore(root_path='./cache')

# 创建CacheBackedEmbeddings
# 在底层Embedding模型的基础上,包装缓存功能
cache_backed_embedding = CacheBackedEmbeddings.from_bytes_store(
    underlying_embeddings=underlying_embedding,  # 指定底层的Embedding模型
    document_embedding_cache=store,  # 指定缓存存储实例
    namespace=underlying_embedding.model,  # 指定缓存的命名空间,用于隔离不同模型的Embedding结果
    query_embedding_cache=True,  # 开启Embedding缓存功能
)

# 使用CacheBackedEmbeddings进行Embedding
text = 'This is a test document.'
vector = cache_backed_embedding.embed_documents([text])
print(f'the dimension of vector is {len(vector[0])}')
print(vector[0][:11])
